// zmq_parallel_tasks.cpp
#include <cstdlib>  // For rand()
#include <iostream>
#include <sstream>
#include <thread>
#include <vector>
#include <zmq.hpp>

#include "serialization.h"
#include "task_result.h"

using namespace Serialization;

/// 序列化 Task 结构体
void encode(std::ostream& os, const Task& task) {
  Serialization::serialize(os, task.row_index);  // 序列化 row_index
  Serialization::serialize(os, task.row);        // 序列化 row
  Serialization::serialize(os, task.B);          // 序列化 B 矩阵
}

/// 反序列化 Task 结构体
void decode(std::istream& is, Task& task) {
  Serialization::deserialize(is, task.row_index);  // 反序列化 row_index
  Serialization::deserialize(is, task.row);        // 反序列化 row
  Serialization::deserialize(is, task.B);          // 反序列化 B 矩阵
}

/// 序列化 Result 结构体
void encode(std::ostream& os, const Result& result) {
  Serialization::serialize(os, result.row_index);   // 序列化 row_index
  Serialization::serialize(os, result.result_row);  // 序列化 result_row
}

/// 反序列化 Result 结构体
void decode(std::istream& is, Result& result) {
  Serialization::deserialize(is, result.row_index);   // 反序列化 row_index
  Serialization::deserialize(is, result.result_row);  // 反序列化 result_row
}

/// Generates a random matrix with the given number of rows and columns.
/// @param rows Number of rows in the matrix.
/// @param cols Number of columns in the matrix.
/// @return A 2D vector representing the matrix.
std::vector<std::vector<std::int32_t>> generateMatrix(uint32_t rows, uint32_t cols) {
  std::vector<std::vector<std::int32_t>> matrix(rows, std::vector<std::int32_t>(cols));
  for (auto& row : matrix) {
    for (auto& elem : row) {
      elem = rand() % 10;  // Random values between 0 and 9
    }
  }
  return matrix;
}

/// Displays the matrix to the standard output.
/// @param matrix The matrix to display.
void displayMatrix(const std::vector<std::vector<std::int32_t>>& matrix) {
  for (const auto& row : matrix) {
    for (const auto& elem : row) {
      std::cout << elem << " ";
    }
    std::cout << std::endl;
  }
}

/// Worker thread function to process tasks.
void workerThread(zmq::context_t& context) {
  try {
    // Socket to receive tasks
    zmq::socket_t task_receiver(context, ZMQ_PULL);
    task_receiver.connect("inproc://tasks");

    // Socket to send results
    zmq::socket_t result_sender(context, ZMQ_PUSH);
    result_sender.connect("inproc://results");

    while (true) {
      // Receive task message
      zmq::message_t task_msg;
      task_receiver.recv(task_msg, zmq::recv_flags::none);

      // Deserialize task
      std::string task_str(static_cast<char*>(task_msg.data()), task_msg.size());

      // Check for STOP signal
      if (task_str == "STOP") {
        break;  // Exit the loop to terminate the thread
      }

      std::istringstream iss(task_str);
      Task task;
      decode(iss, task);

      // Compute the result row
      Result result;
      result.row_index = task.row_index;
      uint32_t cols = static_cast<uint32_t>(task.B[0].size());
      result.result_row.resize(cols, 0);
      for (uint32_t j = 0; j < cols; ++j) {
        for (size_t k = 0; k < task.row.size(); ++k) {
          result.result_row[j] += task.row[k] * task.B[k][j];
        }
      }

      // Serialize the result
      std::ostringstream oss;
      encode(oss, result);
      std::string result_str = oss.str();

      // Send the result message
      zmq::message_t result_msg(result_str.size());
      memcpy(result_msg.data(), result_str.data(), result_str.size());
      result_sender.send(result_msg, zmq::send_flags::none);
    }
  } catch (const zmq::error_t& e) {
    std::cerr << "Worker thread encountered ZMQ error: " << e.what() << std::endl;
  } catch (const std::exception& e) {
    std::cerr << "Worker thread encountered exception: " << e.what() << std::endl;
  }
}

int main() {
  try {
    // Initialize random seed
    std::srand(static_cast<unsigned int>(std::time(nullptr)));

    // Initialize ZMQ context with 1 IO thread
    zmq::context_t context(1);

    // Socket to send tasks
    zmq::socket_t task_sender(context, ZMQ_PUSH);
    task_sender.bind("inproc://tasks");

    // Socket to receive results
    zmq::socket_t result_receiver(context, ZMQ_PULL);
    result_receiver.bind("inproc://results");

    // Launch worker threads
    const uint32_t num_workers = 4;  // Number of worker threads
    std::vector<std::thread> workers;
    for (uint32_t i = 0; i < num_workers; ++i) {
      workers.emplace_back(workerThread, std::ref(context));
    }

    // Generate matrices A and B
    const uint32_t rows = 10;
    const uint32_t cols = 10;
    std::vector<std::vector<std::int32_t>> A = generateMatrix(rows, cols);
    std::vector<std::vector<std::int32_t>> B = generateMatrix(cols, rows);

    std::cout << "Matrix A:" << std::endl;
    displayMatrix(A);
    std::cout << "Matrix B:" << std::endl;
    displayMatrix(B);

    // Send tasks to workers
    for (uint32_t i = 0; i < rows; ++i) {
      Task task;
      task.row_index = i;
      task.row = A[i];
      task.B = B;

      // Serialize the task
      std::ostringstream oss;
      encode(oss, task);
      std::string task_str = oss.str();

      // Send the task message
      zmq::message_t task_msg(task_str.size());
      memcpy(task_msg.data(), task_str.data(), task_str.size());
      task_sender.send(task_msg, zmq::send_flags::none);
    }

    // Send STOP messages to workers to signal them to exit
    for (uint32_t i = 0; i < num_workers; ++i) {
      std::string stop_str = "STOP";
      zmq::message_t stop_msg(stop_str.size());
      memcpy(stop_msg.data(), stop_str.data(), stop_str.size());
      task_sender.send(stop_msg, zmq::send_flags::none);
    }

    // Receive all results
    std::vector<std::vector<std::int32_t>> C(rows, std::vector<std::int32_t>(cols, 0));
    for (uint32_t i = 0; i < rows; ++i) {
      zmq::message_t result_msg;
      result_receiver.recv(result_msg, zmq::recv_flags::none);
      std::string result_str(static_cast<char*>(result_msg.data()), result_msg.size());

      // Deserialize the result
      std::istringstream iss(result_str);
      Result result;
      decode(iss, result);

      // Store the result
      C[result.row_index] = result.result_row;
    }

    std::cout << "Matrix C (Result):" << std::endl;
    displayMatrix(C);

    // Join worker threads to ensure they have terminated gracefully
    for (auto& worker : workers) {
      if (worker.joinable()) {
        worker.join();
      }
    }

    return 0;
  } catch (const zmq::error_t& e) {
    std::cerr << "Producer encountered ZMQ error: " << e.what() << std::endl;
    return 1;
  } catch (const std::exception& e) {
    std::cerr << "Producer encountered exception: " << e.what() << std::endl;
    return 1;
  }
}
